import geopandas as gpd
import numpy as np
import os
from shapely.geometry import box, Polygon, MultiPolygon
from shapely.validation import make_valid
from shapely.ops import unary_union
from shapely.errors import TopologicalError

def generate_all_regions_with_children(geojson_path, output_dir, scale=100000):
    """
    生成所有区域的OBJ模型：
    - 专注于填充空隙，不降低简化阈值
    - 优化顶点匹配但不过度增加计算量
    - 保持子物体结构和导出性能
    """
    # 创建输出目录
    os.makedirs(output_dir, exist_ok=True)
    print(f"✅ 输出目录：{output_dir}（将生成20个区域模型）")

    # 检查输入文件
    if not os.path.exists(geojson_path):
        print(f"❌ 错误：输入文件不存在 - {geojson_path}")
        return

    # ---------------------- 1. 读取并预处理GeoJSON数据 ----------------------
    print("\n" + "="*60)
    print("【步骤1/3】读取并预处理GeoJSON数据")
    print("="*60)
    try:
        gdf = gpd.read_file(geojson_path, dtype={'LIN_CHANGM': str, 'LIN_BAN': str, 'XIAO_BAN': str})
        print(f"✅ 成功读取GeoJSON：共包含 {len(gdf)} 个小地块")

        # 保持原始精度处理
        gdf['geometry'] = gdf['geometry'].apply(lambda geom: round_geometry(geom, 9))

    except Exception as e:
        print(f"❌ 读取GeoJSON失败：{str(e)}")
        return

    # 过滤无效几何并修复
    valid_gdf = gdf[~gdf.geometry.is_empty].copy()
    valid_gdf = valid_gdf[valid_gdf.geometry.type.isin(["Polygon", "MultiPolygon"])].copy()
    valid_gdf['geometry'] = valid_gdf['geometry'].apply(lambda geom: make_valid(geom) if not geom.is_valid else geom)
    valid_gdf = valid_gdf[valid_gdf.geometry.is_valid]

    if len(valid_gdf) == 0:
        print("❌ 错误：无有效小地块几何数据")
        return
    print(f"✅ 过滤并修复后：剩余 {len(valid_gdf)} 个有效小地块")

    # 计算全局经纬度边界
    total_bounds = valid_gdf.total_bounds
    min_lon, min_lat, max_lon, max_lat = total_bounds
    print(f"\n📊 全局经纬度范围：")
    print(f"  经度：{min_lon:.8f} ~ {max_lon:.8f}")
    print(f"  纬度：{min_lat:.8f} ~ {max_lat:.8f}")

    # ---------------------- 2. 生成所有20个区域的边界 ----------------------
    print("\n" + "="*60)
    print("【步骤2/3】生成所有20个区域的边界")
    print("="*60)
    # 计算分割方案
    split_cols, split_rows = calculate_split_grid(20, min_lon, max_lon, min_lat, max_lat)
    lon_step = (max_lon - min_lon) / split_cols
    lat_step = (max_lat - min_lat) / split_rows
    print(f"✅ 分割方案：{split_cols}列 × {split_rows}行 = {split_cols*split_rows}个矩形（取前20个）")

    # 生成所有20个区域的边界信息
    regions = []
    rect_seq = 1  # 区域编号（1-20）

    for col in range(split_cols):
        for row in range(split_rows):
            if rect_seq > 20:
                break  # 仅保留前20个区域

            # 计算当前区域的经纬度边界
            rect_min_lon = min_lon + col * lon_step
            rect_max_lon = min_lon + (col + 1) * lon_step
            rect_min_lat = min_lat + row * lat_step
            rect_max_lat = min_lat + (row + 1) * lat_step

            # 创建矩形几何
            rect_geometry = box(rect_min_lon, rect_min_lat, rect_max_lon, rect_max_lat)
            rect_id = f"Region_{rect_seq:02d}"
            obj_filename = f"{rect_id}_only_plots"

            # 匹配该区域内的小地块
            matched_indices = get_matched_plots(valid_gdf, rect_geometry)
            print(f"  区域 {rect_seq:02d}：匹配到 {len(matched_indices)} 个小地块")

            # 记录区域信息
            regions.append({
                "seq": rect_seq,
                "id": rect_id,
                "filename": obj_filename,
                "bounds": (rect_min_lon, rect_min_lat, rect_max_lon, rect_max_lat),
                "matched_plots": matched_indices
            })

            rect_seq += 1

        if rect_seq > 20:
            break

    # ---------------------- 3. 导出所有区域的OBJ模型（专注填充） ----------------------
    print("\n" + "="*60)
    print("【步骤3/3】导出所有区域的OBJ模型（专注填充空隙）")
    print("="*60)

    # 坐标偏移（以全局中心为原点）
    offset = np.array([(min_lon + max_lon) / 2, (min_lat + max_lat) / 2], dtype=np.float64)
    print(f"✅ 坐标原点（全局中心）：经度={offset[0]:.8f}, 纬度={offset[1]:.8f}")

    # 核心参数（仅调整填充相关参数）
    EXTRUSION_HEIGHT = 0.1
    SIMPLIFY_TOLERANCE = 0.00005  # 保持原始简化阈值，不增加计算量
    MIN_VERTICES = 3
    COORDINATE_PRECISION = 9
    SNAP_THRESHOLD = 0.00008  # 适度提高吸附精度（比原始高但不过度）
    GAP_FILL_TOLERANCE = 0.00015  # 扩大填充范围，专注于空隙填充

    for region in regions:
        seq = region["seq"]
        region_id = region["id"]
        filename = region["filename"]
        matched_plots = region["matched_plots"]
        min_lon_r, min_lat_r, max_lon_r, max_lat_r = region["bounds"]

        output_path = os.path.join(output_dir, f"{filename}.obj")

        try:
            # 收集该区域所有小地块的几何数据
            plot_geometries = []
            # 全局顶点字典，用于跨地块顶点匹配（优化填充的核心）
            global_vertices = {}

            for plot_idx in matched_plots:
                row = valid_gdf.iloc[plot_idx]
                plot_geometry = row.geometry

                # 处理多边形类型
                if isinstance(plot_geometry, MultiPolygon):
                    polygons = [max(plot_geometry.geoms, key=lambda g: g.area)]
                elif isinstance(plot_geometry, Polygon):
                    polygons = [plot_geometry]
                else:
                    continue  # 跳过不支持的几何类型

                if not polygons:
                    continue

                # 小地块命名
                林场 = str(row.get("LIN_CHANGM", "Unknown")).strip().replace(" ", "_").replace("'", "").replace('"', '')
                林班 = str(row.get("LIN_BAN", "0")).strip().zfill(3)
                小班 = str(row.get("XIAO_BAN", "0")).strip().zfill(3)
                plot_unique_id = f"{seq:02d}_{plot_idx:04d}"
                plot_name = f"Plot_{plot_unique_id}_{林场}_{林班}_{小班}"

                # 处理每个子多边形
                for poly_idx, poly in enumerate(polygons):
                    if poly is None or poly.exterior is None:
                        continue

                    # 获取外部轮廓坐标
                    poly_coords = np.array(poly.exterior.coords, dtype=np.float64)
                    if len(poly_coords) < MIN_VERTICES + 1:
                        print(f"⚠️  {region_id}：跳过顶点不足的地块 {plot_name}（{len(poly_coords)}个点）")
                        continue

                    # 处理内部孔洞
                    interiors = []
                    for interior in poly.interiors:
                        if len(interior.coords) >= MIN_VERTICES + 1:
                            interiors.append(np.array(interior.coords, dtype=np.float64))

                    # 多边形简化（保持原始阈值，不增加计算量）
                    simplified_coords = simplify_polygon(poly_coords, tolerance=SIMPLIFY_TOLERANCE)
                    if len(simplified_coords) < MIN_VERTICES + 1:
                        simplified_coords = poly_coords
                        print(f"⚠️  {region_id}：简化后顶点不足，使用原始坐标 {plot_name}")

                    # 增强顶点吸附和空隙填充（仅优化此部分）
                    filled_coords = fill_gaps_and_snap(
                        simplified_coords,
                        global_vertices,
                        SNAP_THRESHOLD,
                        GAP_FILL_TOLERANCE
                    )

                    # 转换为Unity坐标
                    scaled_coords = (filled_coords - offset) * scale

                    # 生成3D顶点（添加高度）
                    vertices = []
                    for coord in scaled_coords:
                        vertices.append([coord[0], coord[1], 0])  # 底部
                        vertices.append([coord[0], coord[1], EXTRUSION_HEIGHT])  # 顶部

                    # 生成面
                    faces = []
                    # 底部面
                    base_vertices = list(range(len(scaled_coords)))
                    for i in range(1, len(base_vertices) - 1):
                        faces.append([base_vertices[0], base_vertices[i], base_vertices[i + 1]])

                    # 顶部面
                    top_vertices = [v + len(scaled_coords) for v in base_vertices]
                    for i in range(len(top_vertices) - 2, 0, -1):
                        faces.append([top_vertices[0], top_vertices[i], top_vertices[i + 1]])

                    # 侧面
                    for i in range(len(base_vertices) - 1):
                        v1 = base_vertices[i]
                        v2 = base_vertices[i + 1]
                        v3 = top_vertices[i + 1]
                        v4 = top_vertices[i]
                        faces.append([v1, v2, v3])
                        faces.append([v1, v3, v4])

                    # 处理内部孔洞（同样应用填充逻辑）
                    hole_vertex_offset = len(vertices)
                    for hole in interiors:
                        simplified_hole = simplify_polygon(hole, tolerance=SIMPLIFY_TOLERANCE/2)
                        if len(simplified_hole) < MIN_VERTICES + 1:
                            continue

                        # 孔洞顶点填充
                        filled_hole = fill_gaps_and_snap(
                            simplified_hole,
                            global_vertices,
                            SNAP_THRESHOLD,
                            GAP_FILL_TOLERANCE
                        )

                        # 添加孔洞顶点
                        hole_vertices = []
                        for coord in filled_hole:
                            scaled = (coord - offset) * scale
                            hole_vertices.append([scaled[0], scaled[1], 0])  # 底部
                            hole_vertices.append([scaled[0], scaled[1], EXTRUSION_HEIGHT])  # 顶部

                        vertices.extend(hole_vertices)

                        # 生成孔洞的面
                        hole_base = list(range(hole_vertex_offset, hole_vertex_offset + len(filled_hole)))
                        hole_top = [v + len(filled_hole) for v in hole_base]

                        # 底部挖空
                        for i in range(len(hole_base) - 2, 0, -1):
                            faces.append([hole_base[0], hole_base[i], hole_base[i + 1]])

                        # 顶部挖空
                        for i in range(1, len(hole_top) - 1):
                            faces.append([hole_top[0], hole_top[i], hole_top[i + 1]])

                        # 孔洞侧面
                        for i in range(len(hole_base) - 1):
                            v1 = hole_base[i]
                            v2 = hole_base[i + 1]
                            v3 = hole_top[i + 1]
                            v4 = hole_top[i]
                            faces.append([v2, v1, v3])
                            faces.append([v3, v1, v4])

                        hole_vertex_offset += len(hole_vertices)

                    # 存储此小地块的几何数据（作为子物体）
                    plot_geometries.append({
                        "name": f"{plot_name}_poly{poly_idx:02d}",
                        "vertices": vertices,
                        "faces": faces
                    })

            # 写入OBJ文件
            with open(output_path, 'w', encoding='utf-8') as f:
                # 1. 写入父对象（区域）
                f.write(f"g {region_id}\n")
                f.write(f"# 区域 {region_id}，包含 {len(plot_geometries)} 个子地块\n")

                # 2. 写入所有顶点
                for plot in plot_geometries:
                    for v in plot["vertices"]:
                        f.write(f"v {v[0]:.{COORDINATE_PRECISION}f} {v[1]:.{COORDINATE_PRECISION}f} {v[2]:.{COORDINATE_PRECISION}f}\n")

                # 3. 写入每个小地块作为独立子物体
                current_offset = 1
                for plot in plot_geometries:
                    f.write(f"\ng {plot['name']}\n")
                    f.write(f"# 子地块：{plot['name']}\n")

                    # 写入面
                    for face in plot["faces"]:
                        f.write(f"f {face[0] + current_offset} {face[1] + current_offset} {face[2] + current_offset}\n")

                    current_offset += len(plot["vertices"])

            progress = (seq / 20) * 100
            print(f"✅ 已生成 {seq}/20 模型：{output_path}（含 {len(plot_geometries)} 个子地块，进度：{progress:.1f}%）")

        except Exception as e:
            print(f"❌ 导出 {region_id} 失败：{str(e)}")
            continue

    print("\n" + "="*60)
    print(f"🎉 所有区域模型生成完成！")
    print(f"📂 模型保存目录：{output_dir}")
    print("💡 提示：已优化填充逻辑，保持简化阈值不变以确保性能")
    print("="*60)


def fill_gaps_and_snap(coords, global_vertices, snap_threshold, gap_fill_threshold):
    """
    专注于空隙填充的顶点处理：
    - 只优化填充逻辑，不增加过多计算
    - 优先填充明显空隙，其次才进行顶点吸附
    """
    # 第一步：填充线性空隙（连续点之间的空隙）
    filled_coords = []
    for i in range(len(coords)):
        current = coords[i]
        filled_coords.append(current)

        # 检查与下一个点之间的距离
        if i < len(coords) - 1:
            next_coord = coords[i+1]
            dist = np.linalg.norm(next_coord - current)

            # 如果距离超过填充阈值，添加中间点填充
            if dist > gap_fill_threshold and dist < gap_fill_threshold * 5:  # 限制过度填充
                # 计算需要添加的点数（根据距离动态调整）
                num_points = max(1, int(dist / gap_fill_threshold))
                for j in range(1, num_points):
                    ratio = j / num_points
                    fill_point = current + (next_coord - current) * ratio
                    filled_coords.append(fill_point)

    # 第二步：顶点吸附（使用更高效的匹配方式）
    snapped = []
    coord_keys = [tuple(np.round(coord, 9)) for coord in filled_coords]

    for i, coord in enumerate(filled_coords):
        key = coord_keys[i]

        # 只检查已有顶点，不进行全局搜索（提高性能）
        if key in global_vertices:
            snapped.append(global_vertices[key])
        else:
            # 只对边界附近的点进行有限吸附
            if i < len(filled_coords) - 1 and i > 0:
                prev_key = coord_keys[i-1]
                next_key = coord_keys[i+1]

                # 检查前后点是否已在全局顶点中
                if prev_key in global_vertices:
                    snapped.append(global_vertices[prev_key])
                    continue
                if next_key in global_vertices:
                    snapped.append(global_vertices[next_key])
                    continue

            # 添加新顶点到全局字典
            snapped.append(coord)
            global_vertices[key] = coord

    return np.array(snapped)


def get_matched_plots(gdf, rect_geometry):
    """获取匹配的地块（提取为独立函数提高可读性）"""
    bbox_mask = gdf.cx[
                rect_geometry.bounds[0]:rect_geometry.bounds[2],
                rect_geometry.bounds[1]:rect_geometry.bounds[3]
                ].index
    candidate_gdf = gdf.loc[bbox_mask].copy()

    # 计算交集面积比例
    candidate_gdf['intersection_area'] = candidate_gdf.geometry.apply(
        lambda geom: geom.intersection(rect_geometry).area / geom.area if geom.area > 0 else 0)
    candidate_gdf['intersection_area'] = candidate_gdf['intersection_area'].astype(float)

    # 获取匹配的地块
    return candidate_gdf[
        (candidate_gdf['intersection_area'] > 0.05) &
        (~np.isnan(candidate_gdf['intersection_area']))
        ].index.tolist()


def calculate_split_grid(split_count, min_lon, max_lon, min_lat, max_lat):
    """计算最优分割网格"""
    lon_range = max_lon - min_lon
    lat_range = max_lat - min_lat
    aspect_ratio = lon_range / lat_range if lat_range != 0 else 2.0

    best_cols = 1
    best_rows = 1
    min_diff = float('inf')

    for cols in range(1, split_count + 1):
        rows = int(np.ceil(split_count / cols))
        total = cols * rows
        current_aspect = (lon_range / cols) / (lat_range / rows) if lat_range != 0 else float('inf')
        aspect_diff = abs(current_aspect - 1.0)
        total_diff = total - split_count

        if (total_diff < min_diff) or (total_diff == min_diff and aspect_diff <
                                       abs(( (lon_range/best_cols)/(lat_range/best_rows) ) - 1.0) if lat_range !=0 else True):
            min_diff = total_diff
            best_cols = cols
            best_rows = rows

    return best_cols, best_rows


def simplify_polygon(points, tolerance):
    """多边形简化算法（保持原始阈值不变）"""
    if len(points) <= 3:
        return points

    def distance(p, a, b):
        """计算点到线段的距离"""
        if np.array_equal(a, b):
            return float(np.linalg.norm(p - a))

        # 转换为三维向量
        a_3d = np.array([a[0], a[1], 0.0])
        b_3d = np.array([b[0], b[1], 0.0])
        p_3d = np.array([p[0], p[1], 0.0])

        cross = np.cross(b_3d - a_3d, p_3d - a_3d)
        return float(np.linalg.norm(cross) / np.linalg.norm(b_3d - a_3d))

    # 找到距离最长的点
    start = points[0]
    end = points[-1]
    max_dist = 0.0
    max_idx = 0

    for i in range(1, len(points)-1):
        dist = distance(points[i], start, end)
        if dist > max_dist:
            max_dist = dist
            max_idx = i

    # 递归简化（使用原始阈值）
    if max_dist > tolerance and not np.isnan(max_dist):
        left = simplify_polygon(points[:max_idx+1], tolerance)
        right = simplify_polygon(points[max_idx:], tolerance)
        return np.vstack((left[:-1], right))
    else:
        return np.array([start, end])


def round_geometry(geom, decimals=9):
    """坐标四舍五入"""
    if geom.is_empty:
        return geom

    def round_coords(coords):
        return [tuple(round(c, decimals) for c in coord) for coord in coords]

    if geom.geom_type == 'Polygon':
        exterior = round_coords(geom.exterior.coords)
        interiors = [round_coords(interior.coords) for interior in geom.interiors]
        return type(geom)(exterior, interiors)
    elif geom.geom_type == 'MultiPolygon':
        polygons = [round_geometry(poly, decimals) for poly in geom.geoms]
        return type(geom)(polygons)
    return geom


# 执行生成
if __name__ == "__main__":
    # 输入输出配置
    input_geojson = r"C:\Users\ma\Desktop\山河屯数据\林场信息.geojson"
    output_directory = r"C:\Users\ma\Desktop\山河屯数据\所有区域带子物体模型"

    # 执行转换
    generate_all_regions_with_children(
        geojson_path=input_geojson,
        output_dir=output_directory,
        scale=100000
    )
